{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Nested Chats for Tool Use in Conversational Chess\n",
    "\n",
    "This notebook demonstrates how to create agents that can play chess with each other\n",
    "while communicating in natural language.\n",
    "The key concept covered in this notebook is the use of nested chats\n",
    "to enable tool use and packaging an LLM-based agent with a tool executor agent\n",
    "into a single agent.\n",
    "\n",
    "Related tutorials:\n",
    "- [Tool Use](https://docs.ag2.ai/latest/docs/user-guide/basic-concepts/introducing-tools/)\n",
    "- [Nested Chats](https://docs.ag2.ai/latest/docs/user-guide/advanced-concepts/orchestration/nested-chat)\n",
    "\n",
    "In this setting, each player is an agent backed by an LLM equipped two tools:\n",
    "- `get_legal_moves` to get a list of current legal moves.\n",
    "- `make_move` to make a move.\n",
    "\n",
    "A board proxy agent is set up to execute the tools and manage the game.\n",
    "It is important to use a board proxy as a non-LLM \"guard rail\" to ensure the game\n",
    "is played correctly and to prevent agents from making illegal moves.\n",
    "\n",
    "Each time a player agent receives a message from the other player agent, \n",
    "it instantiates a nested chat with the board proxy agent to get the legal moves\n",
    "and make a move using the tools given. \n",
    "The nested chat between the player agent and the board agent\n",
    "continues until the a legal move is made by the tool.\n",
    "Once the nested chat concludes, the player agent sends a message to the\n",
    "other player agent about the move made."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Installation\n",
    "\n",
    "First you need to install the `autogen` and `chess` packages to use AG2."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install -qqq autogen chess"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setting up LLMs\n",
    "\n",
    "Now you can set up the models you want to use."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "player_white_config_list = [\n",
    "    {\n",
    "        \"model\": \"gpt-4-turbo-preview\",\n",
    "        \"api_key\": os.environ.get(\"OPENAI_API_KEY\"),\n",
    "    },\n",
    "]\n",
    "\n",
    "player_black_config_list = [\n",
    "    {\n",
    "        \"model\": \"gpt-4-turbo-preview\",\n",
    "        \"api_key\": os.environ.get(\"OPENAI_API_KEY\"),\n",
    "    },\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating tools\n",
    "\n",
    "Write functions for getting legal moves and making a move."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Annotated\n",
    "\n",
    "import chess\n",
    "import chess.svg\n",
    "from IPython.display import display\n",
    "\n",
    "# Initialize the board.\n",
    "board = chess.Board()\n",
    "\n",
    "# Keep track of whether a move has been made.\n",
    "made_move = False\n",
    "\n",
    "\n",
    "def get_legal_moves() -> Annotated[str, \"A list of legal moves in UCI format\"]:\n",
    "    return \"Possible moves are: \" + \",\".join([str(move) for move in board.legal_moves])\n",
    "\n",
    "\n",
    "def make_move(move: Annotated[str, \"A move in UCI format.\"]) -> Annotated[str, \"Result of the move.\"]:\n",
    "    move = chess.Move.from_uci(move)\n",
    "    board.push_uci(str(move))\n",
    "    global made_move\n",
    "    made_move = True\n",
    "    # Display the board.\n",
    "    display(\n",
    "        chess.svg.board(board, arrows=[(move.from_square, move.to_square)], fill={move.from_square: \"gray\"}, size=200)\n",
    "    )\n",
    "    # Get the piece name.\n",
    "    piece = board.piece_at(move.to_square)\n",
    "    piece_symbol = piece.unicode_symbol()\n",
    "    piece_name = (\n",
    "        chess.piece_name(piece.piece_type).capitalize()\n",
    "        if piece_symbol.isupper()\n",
    "        else chess.piece_name(piece.piece_type)\n",
    "    )\n",
    "    return f\"Moved {piece_name} ({piece_symbol}) from {chess.SQUARE_NAMES[move.from_square]} to {chess.SQUARE_NAMES[move.to_square]}.\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating agents\n",
    "\n",
    "Let's create the agents. We have three different agents:\n",
    "- `player_white` is the agent that plays white.\n",
    "- `player_black` is the agent that plays black.\n",
    "- `board_proxy` is the agent that moves the pieces on the board."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from autogen import ConversableAgent, register_function\n",
    "\n",
    "player_white = ConversableAgent(\n",
    "    name=\"Player White\",\n",
    "    system_message=\"You are a chess player and you play as white. \"\n",
    "    \"First call get_legal_moves() first, to get list of legal moves. \"\n",
    "    \"Then call make_move(move) to make a move.\",\n",
    "    llm_config={\"config_list\": player_white_config_list, \"cache_seed\": None},\n",
    ")\n",
    "\n",
    "player_black = ConversableAgent(\n",
    "    name=\"Player Black\",\n",
    "    system_message=\"You are a chess player and you play as black. \"\n",
    "    \"First call get_legal_moves() first, to get list of legal moves. \"\n",
    "    \"Then call make_move(move) to make a move.\",\n",
    "    llm_config={\"config_list\": player_black_config_list, \"cache_seed\": None},\n",
    ")\n",
    "\n",
    "# Check if the player has made a move, and reset the flag if move is made.\n",
    "\n",
    "\n",
    "def check_made_move(msg):\n",
    "    global made_move\n",
    "    if made_move:\n",
    "        made_move = False\n",
    "        return True\n",
    "    else:\n",
    "        return False\n",
    "\n",
    "\n",
    "board_proxy = ConversableAgent(\n",
    "    name=\"Board Proxy\",\n",
    "    llm_config=False,\n",
    "    # The board proxy will only terminate the conversation if the player has made a move.\n",
    "    is_termination_msg=check_made_move,\n",
    "    # The auto reply message is set to keep the player agent retrying until a move is made.\n",
    "    default_auto_reply=\"Please make a move.\",\n",
    "    human_input_mode=\"NEVER\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Register tools for the agents. See [tutorial chapter on tool use](https://docs.ag2.ai/latest/docs/user-guide/basic-concepts/introducing-tools/) \n",
    "for more information."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "register_function(\n",
    "    make_move,\n",
    "    caller=player_white,\n",
    "    executor=board_proxy,\n",
    "    name=\"make_move\",\n",
    "    description=\"Call this tool to make a move.\",\n",
    ")\n",
    "\n",
    "register_function(\n",
    "    get_legal_moves,\n",
    "    caller=player_white,\n",
    "    executor=board_proxy,\n",
    "    name=\"get_legal_moves\",\n",
    "    description=\"Get legal moves.\",\n",
    ")\n",
    "\n",
    "register_function(\n",
    "    make_move,\n",
    "    caller=player_black,\n",
    "    executor=board_proxy,\n",
    "    name=\"make_move\",\n",
    "    description=\"Call this tool to make a move.\",\n",
    ")\n",
    "\n",
    "register_function(\n",
    "    get_legal_moves,\n",
    "    caller=player_black,\n",
    "    executor=board_proxy,\n",
    "    name=\"get_legal_moves\",\n",
    "    description=\"Get legal moves.\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now the agents have their tools ready. You can inspect the auto-generated\n",
    "tool schema for each agent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "player_black.llm_config[\"tools\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Register nested chats for the player agents.\n",
    "Nested chats allows each player agent to chat with the board proxy agent\n",
    "to make a move, before communicating with the other player agent.\n",
    "\n",
    "In the code below, in each nested chat, the board proxy agent starts\n",
    "a conversation with the player agent using the message received from the other\n",
    "player agent (e.g., \"Your move\"). The two agents continue the conversation\n",
    "until a legal move is made using the `make_move` tool.\n",
    "The last message in the nested chat is a message from the player agent about\n",
    "the move made,\n",
    "and this message is then sent to the other player agent.\n",
    "\n",
    "The following diagram illustrates the nested chat between the player agent and the board agent.\n",
    "\n",
    "![Conversational Chess](https://media.githubusercontent.com/media/ag2ai/ag2/main/notebook/nested-chats-chess.png)\n",
    "\n",
    "See [nested chats tutorial chapter](https://docs.ag2.ai/latest/docs/user-guide/advanced-concepts/orchestration/nested-chat)\n",
    "for more information."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "player_white.register_nested_chats(\n",
    "    trigger=player_black,\n",
    "    chat_queue=[\n",
    "        {\n",
    "            # The initial message is the one received by the player agent from\n",
    "            # the other player agent.\n",
    "            \"sender\": board_proxy,\n",
    "            \"recipient\": player_white,\n",
    "            # The final message is sent to the player agent.\n",
    "            \"summary_method\": \"last_msg\",\n",
    "        }\n",
    "    ],\n",
    ")\n",
    "\n",
    "player_black.register_nested_chats(\n",
    "    trigger=player_white,\n",
    "    chat_queue=[\n",
    "        {\n",
    "            # The initial message is the one received by the player agent from\n",
    "            # the other player agent.\n",
    "            \"sender\": board_proxy,\n",
    "            \"recipient\": player_black,\n",
    "            # The final message is sent to the player agent.\n",
    "            \"summary_method\": \"last_msg\",\n",
    "        }\n",
    "    ],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Playing the game\n",
    "\n",
    "Start the chess game."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Clear the board.\n",
    "board = chess.Board()\n",
    "\n",
    "chat_result = player_black.initiate_chat(\n",
    "    player_white,\n",
    "    message=\"Let's play chess! Your move.\",\n",
    "    max_turns=4,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the output above, you can see \"Start a new chat\" is displayed\n",
    "whenever a new nested chat is started between the board proxy agent and a player agent.\n",
    "The \"carryover\" is empty as it is a new chat in the sequence."
   ]
  }
 ],
 "metadata": {
  "front_matter": {
   "description": "LLM-backed agents playing chess with each other using nested chats.",
   "tags": [
    "nested chat",
    "tool/function",
    "orchestration"
   ]
  },
  "kernelspec": {
   "display_name": "autogen",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
